#include <vector>
#include <arm_neon.h>
#include <pthread.h>
#include <iostream>
#include <fstream>
#include <string>
#include <sstream>
#include <cmath>
#include <omp.h>
#include <semaphore.h>
#include <unordered_map>
#include <cstdlib>
#include <sys/time.h>
#define millitime(x) (x.tv_sec * 1000 + x.tv_usec / 1000.0)

using namespace std;

int thread_count;
string paths[4];
string path;
string file_num = "4";
vector<string> file_names;
string file_name;
int COLS, NUMR, NUME;
int totalCols;
int totalRows;
int** A = NULL; 
int* lp = NULL; 
unordered_map<int, int> lpE2R; 

pthread_barrier_t barrier1, barrier2, barrier3, barrier4;

void release_mem() {
    if (A != NULL) {
        for (int i = 0; i < totalRows; i++)
            delete[] A[i];
        delete[] A;
        delete[] lp;
        A = NULL;
        lp = NULL;
    }
}

void initMatrix() {
    A = new int*[totalRows];
    for (int i = 0; i < totalRows; i++) {
        A[i] = new int[totalCols];
        for (int j = 0; j < totalCols; j++)
            A[i][j] = 0;
    }
    lp = new int[totalRows];
    for (int i = 0; i < totalRows; i++)
        lp[i] = -1; // -1表示没有非零元素
    lpE2R.clear();
}

void fresh_arg() {
    string now_file_name = file_name;
    for (int i = 0; i < now_file_name.size(); i++)
        if (now_file_name[i] == '_')
            now_file_name[i] = ' ';
    stringstream ss;
    ss << now_file_name;
    ss >> COLS >> COLS >> NUMR >> NUME;
    totalCols = COLS / 32 + (COLS % 32 != 0 ? 1 : 0);
    totalRows = NUMR + NUME;
}

void readFile(string path, int& nowRow) {
    string line;
    stringstream ss;
    ifstream input(path);
    while (getline(input, line)) {
        int pos;
        ss << line;
        while (ss >> pos) {
            A[nowRow][pos / 32] |= (1 << (pos % 32));
            lp[nowRow] = max(pos, lp[nowRow]);
        }
        if (nowRow >= NUME) {
            if (lpE2R.find(lp[nowRow]) != lpE2R.end()) {
                cout << "Error: Duplicate last non-zero element at row " << lpE2R[lp[nowRow]] << " and " << nowRow - NUME << endl;
                exit(-1);
            }
            lpE2R[lp[nowRow]] = nowRow;
        }
        nowRow++;
        ss.clear();
    }
    input.close();
}

void readData() {
    int nowRow = 0;
    readFile("/openmp/test/" + file_name + "/2.txt", nowRow);
    readFile("/openmp/test/" + file_name + "/1.txt", nowRow);
}

void reset() {
    for (int i = 0; i < totalRows; i++) {
        for (int j = 0; j < totalCols; j++)
            A[i][j] = 0;
    }
    for (int i = 0; i < totalRows; i++)
        lp[i] = -1; // -1表示没有非零元素
    lpE2R.clear();
    readData();
}

void common_gausst() {
    for (int i = 0; i < NUME; i++) {
        while (lp[i] > -1) {
            if (lpE2R.find(lp[i]) != lpE2R.end()) {
                int rowR = lpE2R[lp[i]];
                bool lpHasChanged = false;
                for (int p = 0; p < totalCols; p++) {
                    A[i][p] ^= A[rowR][p];
                }
                for (int p = totalCols - 1; p >= 0; p--) {
                    if (A[i][p] != 0) {
                        lpHasChanged = true;
                        for (int k = 31; k >= 0; k--) {
                            if ((A[i][p] & (1 << k)) != 0) {
                                lp[i] = p * 32 + k;
                                break;
                            }
                        }
                        break;
                    }
                }
                if (!lpHasChanged)
                    lp[i] = -1;
            } else {
                lpE2R[lp[i]] = i;
                break;
            }
        }
    }
}

void* thread_func(void* param) {
    long long t_id = (long long)param;
    for (int i = 0; i < NUME; i++) {
        while (lp[i] > -1) {
            if (lpE2R.find(lp[i]) != lpE2R.end()) {
                int rowR = lpE2R[lp[i]];
                for (int p = t_id; p < totalCols; p += thread_count) {
                    A[i][p] ^= A[rowR][p];
                }
                pthread_barrier_wait(&barrier1);
                if (t_id == 0) {
                    bool lpHasChanged = false;
                    for (int p = totalCols - 1; p >= 0; p--) {
                        if (A[i][p] != 0) {
                            lpHasChanged = true;
                            for (int k = 31; k >= 0; k--) {
                                if ((A[i][p] & (1 << k)) != 0) {
                                    lp[i] = p * 32 + k;
                                    break;
                                }
                            }
                            break;
                        }
                    }
                    if (!lpHasChanged)
                        lp[i] = -1;
                }
                pthread_barrier_wait(&barrier2);
            } else {
                pthread_barrier_wait(&barrier4);
                if (t_id == 0)
                    lpE2R[lp[i]] = i;
                pthread_barrier_wait(&barrier3);
                break;
            }
        }
    }
    return NULL;
}

void static_gausst_p() {
    pthread_barrier_init(&barrier1, NULL, thread_count);
    pthread_barrier_init(&barrier2, NULL, thread_count);
    pthread_barrier_init(&barrier3, NULL, thread_count);
    pthread_barrier_init(&barrier4, NULL, thread_count);
    pthread_t* handles = new pthread_t[thread_count];
    for (int i = 0; i < thread_count; i++)
        pthread_create(handles + i, NULL, thread_func, (void*)(long long)i);
    for (int i = 0; i < thread_count; i++)
        pthread_join(handles[i], NULL);
    delete[] handles;
    pthread_barrier_destroy(&barrier1);
    pthread_barrier_destroy(&barrier2);
    pthread_barrier_destroy(&barrier3);
    pthread_barrier_destroy(&barrier4);
}

void open_mp_gausst() {
    #pragma omp parallel num_threads(thread_count)
    for (int i = 0; i < NUME; i++) {
        while (lp[i] > -1) {
            if (lpE2R.find(lp[i]) != lpE2R.end()) {
                int rowR = lpE2R[lp[i]];
                #pragma omp for
                for (int p = 0; p < totalCols; p++) {
                    A[i][p] ^= A[rowR][p];
                }
                #pragma omp master
                {
                    bool lpHasChanged = false;
                    for (int p = totalCols - 1; p >= 0; p--) {
                        if (A[i][p] != 0) {
                            lpHasChanged = true;
                            for (int k = 31; k >= 0; k--) {
                                if ((A[i][p] & (1 << k)) != 0) {
                                    lp[i] = p * 32 + k;
                                    break;
                                }
                            }
                            break;
                        }
                    }
                    if (!lpHasChanged)
                        lp[i] = -1;
                }
                #pragma omp barrier
            } else {
                #pragma omp barrier
                #pragma omp master
                lpE2R[lp[i]] = i;
                #pragma omp barrier
                break;
            }
        }
    }
}

void open_mp_gausst_simd() {
    #pragma omp parallel num_threads(thread_count)
    for (int i = 0; i < NUME; i++) {
        while (lp[i] > -1) {
            if (lpE2R.find(lp[i]) != lpE2R.end()) {
                int rowR = lpE2R[lp[i]];
                #pragma omp for simd safelen(4)
                for (int p = 0; p < totalCols; p++) {
                    A[i][p] ^= A[rowR][p];
                }
                #pragma omp master
                {
                    bool lpHasChanged = false;
                    for (int p = totalCols - 1; p >= 0; p--) {
                        if (A[i][p] != 0) {
                            lpHasChanged = true;
                            for (int k = 31; k >= 0; k--) {
                                if ((A[i][p] & (1 << k)) != 0) {
                                    lp[i] = p * 32 + k;
                                    break;
                                }
                            }
                            break;
                        }
                    }
                    if (!lpHasChanged)
                        lp[i] = -1;
                }
                #pragma omp barrier
            } else {
                #pragma omp barrier
                #pragma omp master
                lpE2R[lp[i]] = i;
                #pragma omp barrier
                break;
            }
        }
    }
}

int main(int argc, char* argv[]) {
    struct timeval head, tail;
    if (argc != 3) {
        cout << "usage: ./gausst_p.exe file_name thread_count" << endl;
        return -1;
    }
    thread_count = atoi(argv[2]);
    file_name = argv[1];
    fresh_arg();
    initMatrix();
    readData();
    gettimeofday(&head, NULL);
    common_gausst();
    gettimeofday(&tail, NULL);
    cout << "Common time: " << millitime(tail) - millitime(head) << "ms" << endl;

    reset();
    gettimeofday(&head, NULL);
    static_gausst_p();
    gettimeofday(&tail, NULL);
    cout << "Pthread time: " << millitime(tail) - millitime(head) << "ms" << endl;

    reset();
    gettimeofday(&head, NULL);
    open_mp_gausst();
    gettimeofday(&tail, NULL);
    cout << "OpenMP time: " << millitime(tail) - millitime(head) << "ms" << endl;

    reset();
    gettimeofday(&head, NULL);
    open_mp_gausst_simd();
    gettimeofday(&tail, NULL);
    cout << "OpenMP SIMD time: " << millitime(tail) - millitime(head) << "ms" << endl;

    release_mem();
    return 0;
}
